import torch
from torch import nn
from fasttext import FTModel
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv, GCNConv
from torch_geometric.nn import GATConv, Linear, to_hetero
from tokenizer_tools import bert_model, bert_tokenizer
from create_alarms_data_002 import ALL_ORIGIN_INFOS


class Classifier(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(Classifier, self).__init__()
        self.project = nn.Sequential(
            # nn.BatchNorm1d(hidden_size),
            nn.Linear(input_size, hidden_size, bias=False),
            # nn.LayerNorm(hidden_size),
            nn.Tanh(),
            # nn.LeakyReLU(0.2),
            # nn.LayerNorm(),
            nn.Linear(hidden_size, hidden_size,bias=False),
            # nn.BatchNorm1d(hidden_size)
        )

    def forward(self, x, index_edge_label):
        # print(x.shape)
        # print(index_edge_label.shape)
        x = self.project(x)
        edge_feat_A = x[index_edge_label[0]]
        edge_feat_B = x[index_edge_label[1]]
        out = (edge_feat_A * edge_feat_B).sum(dim=-1)
        out = torch.sigmoid(out)
        return out


class GAT_Model(nn.Module):
    def __init__(self, hidden_channels):
        super(GAT_Model, self).__init__()
        self.conv1 = GATConv(hidden_channels, hidden_channels, heads=4)
        self.conv2 = GATConv(hidden_channels * 4, hidden_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


class GNN(nn.Module):
    def __init__(self, hidden_channels):
        super(GNN, self).__init__()
        # self.BN = nn.BatchNorm1d(300)
        self.conv1 = SAGEConv(hidden_channels, 32)
        self.conv2 = SAGEConv(32, 16)
        # self.conv3 =GCNConv(hidden_channels, hidden_channels)
        # self.conv3 = GATConv(hidden_channels, hidden_channels)
        # self.conv4 = GATConv(hidden_channels, hidden_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return x


class Model(nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.gnn = GAT_Model(hidden_channels)
        self.classifier = Classifier(hidden_channels, hidden_channels)

    def change_model(self, data):
        self.gnn = to_hetero(self.gnn, metadata=data.metadata())

    def forward(self, batch_data):
        new_x = self.gnn(batch_data.x, batch_data.edge_label_index)
        out = self.classifier(new_x, batch_data.edge_label_index)
        return out
